import { PlusOutlined } from "@ant-design/icons";
import { ActionType, ProColumns, ProTable } from "@ant-design/pro-components";
import { Button, Form, Input, message, Modal, Select } from "antd"
import { memo, useRef, useState } from "react";
import { getDataSetList, training, trainingLog, tuningCreate, tuningDelete, tuningList, tuningUpdate } from "../api/request";
import { formatTime } from "../utils/utils";
import CustomSelect from "../components/CustomSelect";
import { useRequest } from "ahooks";

const layout = {
    labelCol: { span: 10 },
    wrapperCol: { span: 12 },
};

// 配置任务
const FineTuning = () => {
    const actionRef = useRef<ActionType>();
    const [visibleCreate, setvisibleCreate] = useState(false);
    const [type, setType] = useState<'create' | 'update'>('create');
    const [form] = Form.useForm();
    const [llmTuningData, setLlmTuningData] = useState<LlmTuning>(null!)
    const [loadding, setLoading] = useState<boolean>(null!)
    const formRef = useRef(null!)
    const { data: trainingLogData, run: trainingLogRun, loading: Loading } = useRequest(trainingLog, {
        manual: true
    })
    const [trainingLogOpen, setTrainingLogOpen] = useState<boolean>(false)

    const config = ({ name,
        model_name = "Qwen2-0.5B-Instruct",
        num_train_epochs = 3,  // 训练轮数
        learning_rate = "3e-05", // 学习率
        max_samples = 100000, // 最大样本数
        cutoff_len = 1024,   // 截断长度
        save_steps = 100,  // 每100步评估一次模型性能
    }) => {

        return `llamafactory-cli train \
--stage sft \
--do_train True \
--model_name_or_path /app/qwen/${model_name} \
--preprocessing_num_workers 16 \
--finetuning_type lora \
--template default \
--flash_attn auto \
--dataset_dir data \
--dataset data_set \
--cutoff_len ${cutoff_len} \
--learning_rate ${learning_rate} \
--num_train_epochs ${num_train_epochs} \
--max_samples ${max_samples} \
--per_device_train_batch_size 2 \
--gradient_accumulation_steps 8 \
--lr_scheduler_type cosine \
--max_grad_norm 1.0 \
--logging_steps 5 \
--save_steps ${save_steps} \
--warmup_steps 0 \
--optim adamw_torch \
--packing False \
--report_to none \
--output_dir /app/saves/Qwen2-0.5B/lora/${name} \
--fp16 True \
--plot_loss True \
--ddp_timeout 180000000 \
--include_num_input_tokens_seen True \
--lora_rank 8 \
--lora_alpha 16 \
--lora_dropout 0 \
--lora_target all
`}

    const columns: ProColumns<LlmTuning>[] = [
        {
            title: '名称',
            dataIndex: 'name',
            width: 200,
        },
        {
            title: 'ROLA微调配置',
            dataIndex: 'config',
            search: false,
            render: (text, record,) => {
                return <div
                    style={
                        {
                            whiteSpace: 'pre-wrap'
                        }
                    }>{record.config}</div>
            }
        },
        {
            title: '状态',
            width: 180,
            dataIndex: 'status',
            search: false,
            valueEnum: {
                close: { text: '未开始', status: 'Default' },
                running: { text: '运行中', status: 'Processing' },
                online: { text: '完成', status: 'Success' },
                error: { text: '异常', status: 'Error' },
            },
        },
        {
            title: '创建时间',
            dataIndex: 'createdTime',
            search: false,
            width: 230,
            render: (text, record,) => {
                return formatTime(record.createdTime)
            }
        },
        {
            title: '操作',
            valueType: 'option',
            key: 'option',
            width: 250,
            render: (text, record, _, action) => [
                <Button type="link" style={{ padding: 0 }} onClick={() => {
                    setType('update')
                    setvisibleCreate(true)
                    setLlmTuningData(record)
                    form.setFieldsValue(record)
                }}>
                    修改
                </Button>,
                <Button type="link" style={{ padding: 0 }} onClick={() => {
                    Modal.confirm({
                        title: '删除',
                        content: '是否删除',
                        onOk: async () => {
                            await tuningDelete(record.id)
                            actionRef.current?.reload()
                            Modal.destroyAll();
                            message.success("删除成功")
                        }
                    })
                }}>
                    删除
                </Button>,
                <Button type="link" style={{ padding: 0 }} onClick={() => {
                    setTrainingLogOpen(true)
                    trainingLogRun(record.id)
                }}>
                    训练日志
                </Button>,
                <Button type="link" style={{ padding: 0 }} onClick={() => {
                    Modal.confirm({
                        title: '训练',
                        content: '是否开始训练模型',
                        onOk: async () => {
                            await training(record.id)
                            message.success('开始训练')
                            actionRef.current?.reload()
                        }
                    })
                }}>
                    训练
                </Button>,
                <Button type="link" style={{ padding: 0 }} href="http://10.10.40.102:17860/" target="_blank">
                    模型测试
                </Button>
            ],
        }
    ]

    return <div>
        <ProTable<LlmTuning>
            columns={columns}
            actionRef={actionRef}
            headerTitle="配置任务"
            cardBordered
            editable={{
                type: 'multiple',
            }}
            formRef={formRef}
            rowKey="id"
            options={{
                setting: {
                    listsHeight: 400,
                },
            }}
            pagination={{
                pageSize: 5
            }}
            request={async (params) => {
                return await tuningList(params.name)
            }}
            dateFormatter="string"
            toolBarRender={() => [
                <Button
                    key="button"
                    icon={<PlusOutlined />}
                    onClick={() => (setvisibleCreate(true), setType("create"))}
                    type="primary"
                >
                    新建
                </Button>
            ]}
        />

        <Modal title={type == "create" ? "新建" : "修改"}
            open={visibleCreate}
            confirmLoading={loadding}
            onOk={() => {
                form.validateFields().then(async (values) => {
                    const configData = config(values)
                    setLoading(true)
                    if (type == "create") {
                        const result = await tuningCreate({
                            ...values,
                            config: configData
                        })
                        if (result.code == 200) {
                            message.success('创建成功')
                            setLoading(false)
                        }
                        console.log(result)
                    } else {
                        const result = await tuningUpdate({
                            id: llmTuningData.id, ...values,
                            config: configData
                        })
                        if (result.code == 200) {
                            message.success('修改成功')
                            setLoading(false)
                        }
                    }
                    actionRef.current?.reload()
                    setvisibleCreate(false)
                })
            }}
            onCancel={() => setvisibleCreate(false)}
            onClose={() => setvisibleCreate(false)}>
            <Form
                {...layout}
                form={form}
                style={{ maxWidth: 600 }}
                className="mt-6"
            >
                <Form.Item name="name" label="任务名称" rules={[{ required: true }]}>
                    <Input autoComplete="off" placeholder="任务名称" />
                </Form.Item>
                <Form.Item name="modelName" label="微调基座模型" rules={[{ required: true }]}>
                    <Select options={[{
                        label: "Qwen2-0.5B-Instruct",
                        value: "Qwen2-0.5B-Instruct"
                    },
                    {
                        label: "Qwen2-7B-Instruct",
                        value: "Qwen2-7B-Instruct"
                    },
                    {
                        label: "Qwen2-14B-Instruct",
                        value: "Qwen2-14B-Instruct"
                    }
                    ]} placeholder="模型选择" />
                </Form.Item>
                <Form.Item name="dataSet" label="数据集" rules={[{ required: true }]}>
                    <CustomSelect api={getDataSetList} fieldNames={{ label: "name", value: "file_url" }} />
                </Form.Item>
                <Form.Item name="cutoff_len" label="截断长度(cutoff_len)">
                    <Input autoComplete="off" placeholder="1024" />
                </Form.Item>
                <Form.Item name="max_samples" label="最大样本数(max_samples)">
                    <Input autoComplete="off" placeholder="100" />
                </Form.Item>
                <Form.Item name="learning_rate" label="学习率(learning_rate)">
                    <Input autoComplete="off" placeholder="3e-05" />
                </Form.Item>
                <Form.Item name="num_train_epochs" label="训练轮数(num_train_epochs)">
                    <Input autoComplete="off" placeholder="3" />
                </Form.Item>
                <Form.Item name="save_steps" label="训练步数(save_steps)">
                    <Input autoComplete="off" placeholder="100" />
                </Form.Item>
            </Form>
        </Modal>

        <Modal title="训练日志" width={900} open={trainingLogOpen} loading={Loading} onCancel={() => setTrainingLogOpen(false)} footer={null}>
            <div style={{ whiteSpace: "pre-wrap", height: 500, overflow: "auto", border: '1px solid #D9D9D9', borderRadius: '10px' }} >
                {trainingLogData?.data as string}
            </div>
        </Modal>
    </div>
}

export default memo(FineTuning)